import os
import numpy as np
from PIL import Image
from multiprocessing import Pool

data = './data/neu_seg_competition/annotations/test'


def unique_pixel_values(image_path):
    # 打开图片
    image = Image.open(image_path)
    # 转换为numpy数组
    image_array = np.array(image)
    # 获取唯一像素值
    unique_values = np.unique(image_array)
    return unique_values


def process_images(image_paths):
    # 创建进程池
    num_threads = os.cpu_count()
    with Pool(processes=num_threads) as pool:
        # 使用map函数并行处理
        results = pool.imap(unique_pixel_values, image_paths)

    # 合并所有结果
    combined_result = np.concatenate(results)
    # 获取整体唯一值
    unique_values = np.unique(combined_result)

    return unique_values

if __name__ == '__main__':
    image_files = [os.path.join(data, filename) for filename in os.listdir(data) if filename.endswith('.png')]
    unique_values = process_images(image_files)
    print(unique_values)